from torch import nn
import collections
import math
import torch
from d2l import torch as d2l

import ai


# @save
class Encoder(nn.Module):
    """编码器-解码器架构的基本编码器接⼝"""

    def __init__(self, **kwargs):
        super(Encoder, self).__init__(**kwargs)

    def forward(self, X, *args):
        raise NotImplementedError


# @save
class Decoder(nn.Module):
    """编码器-解码器架构的基本解码器接⼝"""

    def __init__(self, **kwargs):
        super(Decoder, self).__init__(**kwargs)

    def init_state(self, enc_outputs, *args):
        raise NotImplementedError

    def forward(self, X, state):
        raise NotImplementedError


# @save
class EncoderDecoder(nn.Module):
    """编码器-解码器架构的基类"""

    def __init__(self, encoder, decoder, **kwargs):
        super(EncoderDecoder, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, enc_X, dec_X, *args):
        enc_outputs = self.encoder(enc_X, *args)
        dec_state = self.decoder.init_state(enc_outputs, *args)
        return self.decoder(dec_X, dec_state)


# @save
class Seq2SeqEncoder(Encoder):
    """⽤于序列到序列学习的循环神经⽹络编码器"""

    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 dropout=0, **kwargs):
        super(Seq2SeqEncoder, self).__init__(**kwargs)
        # 嵌⼊层
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size, num_hiddens, num_layers,
                          dropout=dropout)

    def forward(self, X, *args):
        # 输出'X'的形状：(batch_size,num_steps,embed_size)
        X = self.embedding(X)
        # 在循环神经⽹络模型中，第⼀个轴对应于时间步
        X = X.permute(1, 0, 2)
        # 如果未提及状态，则默认为0
        output, state = self.rnn(X)
        # output的形状:(num_steps,batch_size,num_hiddens)
        # state[0]的形状:(num_layers,batch_size,num_hiddens)
        return output, state


class Seq2SeqDecoder(d2l.Decoder):
    """⽤于序列到序列学习的循环神经⽹络解码器"""

    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 dropout=0, **kwargs):
        super(Seq2SeqDecoder, self).__init__(**kwargs)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(embed_size + num_hiddens, num_hiddens, num_layers,
                          dropout=dropout)
        self.dense = nn.Linear(num_hiddens, vocab_size)

    def init_state(self, enc_outputs, *args):
        return enc_outputs[1]

    def forward(self, X, state):
        # 输出'X'的形状：(batch_size,num_steps,embed_size)
        X = self.embedding(X).permute(1, 0, 2)
        # ⼴播context，使其具有与X相同的num_steps
        context = state[-1].repeat(X.shape[0], 1, 1)  # 使用encode的最后一个step的状态作为上下文,重复X_steps次分别拼接至X
        X_and_context = torch.cat((X, context), 2)
        output, state = self.rnn(X_and_context, state)
        output = self.dense(output).permute(1, 0, 2)  # output中每一步都对应一个预测值,组成一个输出序列,每个预测值参考的状态都是在其前面的
        # output的形状:(batch_size,num_steps,vocab_size)
        # state[0]的形状:(num_layers,batch_size,num_hiddens)
        return output, state


def sequence_mask(X, valid_len, value=0):
    """在序列中屏蔽不相关的项"""
    maxlen = X.size(1)
    mask = torch.arange((maxlen), dtype=torch.float32,
                        device=X.device)[None, :] < valid_len[:, None]
    X[~mask] = value
    return X


# @save
class MaskedSoftmaxCELoss(nn.CrossEntropyLoss):
    """带遮蔽的softmax交叉熵损失函数"""

    # pred的形状：(batch_size,num_steps,vocab_size)
    # label的形状：(batch_size,num_steps)
    # valid_len的形状：(batch_size,)
    def forward(self, pred, label, valid_len):
        weights = torch.ones_like(label)
        weights = sequence_mask(weights, valid_len)
        self.reduction = 'none'
        unweighted_loss = super(MaskedSoftmaxCELoss, self).forward(
            pred.permute(0, 2, 1), label)  # 规定需要(batch_size,vocab_size,step_nums)
        weighted_loss = (unweighted_loss * weights).mean(dim=1)
        return weighted_loss


from ai import AiConstant
from project.enchanter_gpt import tokenize_vocab
from project.enchanter_gpt.tokenize_vocab import gxl_tokenizer

import os
from tqdm import tqdm


def read_txt_file(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        content = file.read()
    content = content.replace("\n", "")
    content = content.replace("。", gxl_tokenizer.EOD_T + "。" + gxl_tokenizer.SOD_T)
    content = gxl_tokenizer.SOD_T + content + gxl_tokenizer.EOD_T

    return content


def process_text(text, chunk_size=20):
    chunks = [text[i:i + chunk_size] for i in range(0, len(text), 1) if len(text[i:i + chunk_size]) == chunk_size]
    return chunks


def load_datalist_from_disk(directory_path=AiConstant.DATA_PATH + "gpt/data/作文数据集/",
                            tokenizer=tokenize_vocab.gxl_tokenizer):
    txt_files = [f for f in os.listdir(directory_path) if f.endswith('.txt')]
    data_list = []
    a = 0
    for file_name in tqdm(txt_files):
        a = a + 1
        if a > 1200:
            break
        file_path = os.path.join(directory_path, file_name)
        content = read_txt_file(file_path)
        chunks = process_text(content)
        data_list = data_list + chunks
    data_list = [tokenizer.encoder([a for a in line]) for line in data_list]
    return data_list


from torch.utils.data import Dataset, DataLoader


class ChatDataSet(Dataset):
    def __init__(self, data_list):
        self.data = data_list

    def __len__(self):
        return len(self.data) - 20

    def __getitem__(self, index):
        return torch.tensor(self.data[index]), torch.tensor(self.data[index + 20])


def get_dataloader():
    data_list = load_datalist_from_disk()
    return DataLoader(ChatDataSet(data_list), batch_size=1000)


class UnMaskedSoftmaxCELoss(nn.CrossEntropyLoss):
    """带遮蔽的softmax交叉熵损失函数"""

    # pred的形状：(batch_size,num_steps,vocab_size)
    # label的形状：(batch_size,num_steps)
    # valid_len的形状：(batch_size,)
    def forward(self, pred, label, valid_len):
        # weights = torch.ones_like(label)
        # weights = sequence_mask(weights, valid_len)
        self.reduction = 'none'
        unweighted_loss = super(MaskedSoftmaxCELoss, self).forward(
            pred.permute(0, 2, 1), label)  # 规定需要(batch_size,vocab_size,step_nums)
        weighted_loss = (unweighted_loss).mean(dim=1)
        return weighted_loss


def train_seq2seq(net, data_iter, lr, num_epochs, device):
    """训练序列到序列模型"""

    def xavier_init_weights(m):
        if type(m) == nn.Linear:
            nn.init.xavier_uniform_(m.weight)
        if type(m) == nn.GRU:
            for param in m._flat_weights_names:
                if "weight" in param:
                    nn.init.xavier_uniform_(m._parameters[param])

    net.apply(xavier_init_weights)
    net.to(device)
    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
    loss = UnMaskedSoftmaxCELoss()
    net.train()
    animator = d2l.Animator(xlabel='epoch', ylabel='loss', xlim=[10, num_epochs])
    for epoch in range(num_epochs):
        timer = d2l.Timer()
        metric = d2l.Accumulator(2)  # 训练损失总和，词元数量
        for batch in data_iter:
            optimizer.zero_grad()
            X, Y = [x.to(device) for x in batch]
            # bos = torch.tensor([tgt_vocab['<bos>']] * Y.shape[0],
            #                    device=device).reshape(-1, 1)
            # dec_input = torch.cat([bos, Y[:, :-1]], 1)  # 强制教学, 强制让模型输出的句子以bos开头
            Y_hat, _ = net(X, Y)
            l = loss(Y_hat, Y, 20)
            l.sum().backward()  # 损失函数的标量进⾏“反向传播”
            d2l.grad_clipping(net, 1)
            # num_tokens = Y_valid_len.sum()
            optimizer.step()
            with torch.no_grad():
                metric.add(l.sum(), 1)
        if (epoch + 1) % 1 == 0:
            animator.add(epoch + 1, (metric[0] / metric[1],))
    # print(f'loss {metric[0] / metric[1]:.3f}, {metric[1] / timer.stop():.1f} tokens/sec on {str(device)}')
    torch.save(net.state_dict(), ai.AiConstant.OUTPUT_PATH + 'gpt/model_params.pth')


def train_run():
    embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.1
    batch_size, num_steps = 64, 10
    lr, num_epochs, device = 0.005, 10, torch.device('cuda:2')
    train_iter = get_dataloader()
    encoder = Seq2SeqEncoder(tokenize_vocab.gxl_tokenizer.vocab_size, embed_size, num_hiddens, num_layers,
                             dropout)
    decoder = Seq2SeqDecoder(tokenize_vocab.gxl_tokenizer.vocab_size, embed_size, num_hiddens, num_layers,
                             dropout)
    net = EncoderDecoder(encoder, decoder)
    train_seq2seq(net, train_iter, lr, num_epochs, device)


from project.enchanter_gpt.tokenize_vocab import gxl_tokenizer


def predict_seq2seq2(net, src_sentence, num_steps,
                     device=torch.device('cpu'), save_attention_weights=False):
    """序列到序列模型的预测"""
    # 在预测时将net设置为评估模式
    net.eval()
    net = net.to(device)
    src_tokens = torch.tensor(gxl_tokenizer.encoder([i for i in src_sentence]))
    print(src_tokens)
    # 添加批量轴
    enc_X = torch.unsqueeze(
        torch.tensor(src_tokens, dtype=torch.long, device=device), dim=0)
    enc_outputs = net.encoder(enc_X)
    dec_state = net.decoder.init_state(enc_outputs)
    # 添加批量轴
    dec_X = torch.unsqueeze(torch.tensor(
        [gxl_tokenizer.encoder('我')], dtype=torch.long, device=device), dim=0)
    output_seq, attention_weight_seq = [], []
    for _ in range(num_steps):
        Y, dec_state = net.decoder(dec_X, dec_state)
        # 我们使⽤具有预测最⾼可能性的词元，作为解码器在下⼀时间步的输⼊
        dec_X = Y.argmax(dim=2)
        pred = dec_X.squeeze(dim=0).TYPE(torch.int32).item()
        # 保存注意⼒权重（稍后讨论）
        if save_attention_weights:
            attention_weight_seq.append(net.decoder.attention_weights)
        # ⼀旦序列结束词元被预测，输出序列的⽣成就完成了
        if pred == gxl_tokenizer.encoder(']'):
            break
        output_seq.append(pred)
    return (gxl_tokenizer.decoder(output_seq)), attention_weight_seq


embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.1
batch_size, num_steps = 64, 10
lr, num_epochs, device = 0.005, 50, torch.device('cuda:1')
train_iter = get_dataloader()
encoder = Seq2SeqEncoder(tokenize_vocab.gxl_tokenizer.vocab_size, embed_size, num_hiddens, num_layers,
                         dropout)
decoder = Seq2SeqDecoder(tokenize_vocab.gxl_tokenizer.vocab_size, embed_size, num_hiddens, num_layers,
                         dropout)
net = EncoderDecoder(encoder, decoder)
print(tokenize_vocab.gxl_tokenizer.vocab_size)
net.load_state_dict(torch.load(ai.AiConstant.OUTPUT_PATH + 'gpt/model_params.pth'))
print(predict_seq2seq2(net, '你好啊，我是雪龙，你是谁呀？', 10))
